// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/parquet/reader/string_column_reader.h>

#include <pollux/dwio/common/selective_column_reader_internal.h>

namespace kumo::pollux::parquet {
    StringColumnReader::StringColumnReader(
        const std::shared_ptr<const dwio::common::TypeWithId> &fileType,
        ParquetParams &params,
        common::ScanSpec &scanSpec)
        : SelectiveColumnReader(fileType->type(), fileType, params, scanSpec) {
    }

    uint64_t StringColumnReader::skip(uint64_t numValues) {
        formatData_->skip(numValues);
        return numValues;
    }

    void StringColumnReader::read(
        int64_t offset,
        const RowSet &rows,
        const uint64_t *incomingNulls) {
        prepareRead<melon::StringPiece>(offset, rows, incomingNulls);
        dwio::common::StringColumnReadWithVisitorHelper<true, false>(
            *this, rows)([&](auto visitor) {
            formatData_->as<ParquetData>().readWithVisitor(visitor);
        });
        readOffset_ += rows.back() + 1;
    }

    void StringColumnReader::getValues(const RowSet &rows, VectorPtr *result) {
        if (scanState_.dictionary.values) {
            auto dictionaryValues =
                    formatData_->as<ParquetData>().dictionaryValues(fileType_->type());
            compactScalarValues<int32_t, int32_t>(rows, false);

            *result = std::make_shared<DictionaryVector<StringView> >(
                memoryPool_, resultNulls(), numValues_, dictionaryValues, values_);
            return;
        }
        rawStringBuffer_ = nullptr;
        rawStringSize_ = 0;
        rawStringUsed_ = 0;
        getFlatValues<StringView, StringView>(rows, result, fileType_->type());
    }

    void StringColumnReader::dedictionarize() {
        if (scanSpec_->keepValues()) {
            auto dict = formatData_->as<ParquetData>()
                    .dictionaryValues(fileType_->type())
                    ->as<FlatVector<StringView> >();
            auto valuesCapacity = values_->capacity();
            auto indices = values_->as<vector_size_t>();
            // 'values_' is sized for the batch worth of StringViews. It is filled with
            // 32 bit indices by dictionary scan.
            POLLUX_CHECK_GE(valuesCapacity, numValues_ * sizeof(StringView));
            stringBuffers_.clear();
            rawStringBuffer_ = nullptr;
            rawStringSize_ = 0;
            rawStringUsed_ = 0;
            auto numValues = numValues_;
            // Convert indices to values in place. Loop from end to beginning
            // so as not to overwrite integer indices with longer StringViews.
            for (auto i = numValues - 1; i >= 0; --i) {
                if (anyNulls_ && bits::isBitNull(rawResultNulls_, i)) {
                    reinterpret_cast<StringView *>(rawValues_)[i] = StringView();
                    continue;
                }
                auto &view = dict->value_at(indices[i]);
                numValues_ = i;
                addStringValue(melon::StringPiece(view.data(), view.size()));
            }
            numValues_ = numValues;
        }
        scanState_.clear();
        formatData_->as<ParquetData>().clearDictionary();
    }
} // namespace kumo::pollux::parquet
